#include "StdAfx.h"
#include "ExpandTimeSteps.h"
#include <iostream>
#include <iosfwd>
#include <sstream>

//////////////////////////////////////////////////////////////////////////
//Author	:	Ross Conroy ross.conroy@tees.ac.uk
//Date		:	06/05/2014
//
//The purpose of this class is to expand an influence diagrams time steps,
//the starting ID requires at least two time steps in order to determine
//how nodes in each time step are linked, names of nodes in ID are used to
//determine what time step they belong to and must have the format
//<node type><node name><time step> e.g. dAction0 for a decsion node in
//the first time step.
//
//Algorithm
//	1.Find last time step in ID, this can be done by finding the nodes
//		with the highest ending number (in most cases 1 or 2).
//	2.Loop through all nodes that belong to the time step and determine
//		what nodes and time-step they link to.
//	3.Generate new nodes and link them into the ID
//		3.1 Create new nodes with states/actions and add to List of new 
//			nodes.
//			3.1.1 Use previous time step to determine what the parent nodes
//			will be for next time step.
//				(IF DID) - decision nodes need to be linked to all previous
//				action and decision nodes, to link action nodes add all
//				action nodes that have a lower time step. To link
//				observation nodes use previous time step to get observation
//				node name and add it and all previous time step nodes.
//		3.2 Loop through List of new nodes and check whether they can be
//			added to the ID, they can be added when all parent nodes are
//			available in the ID. If a node cannot be added move onto the
//			next node and go back to the beginning of the List once a node
//			has been successfully added.
//	4.Once a time step has been added go back to previous time step to get
//		the parent order for each node and apply it to the new nodes.
//	5.Copy CPT's to new nodes from previous time step for all chance nodes
//	6.Repeat for number of time steps to be added.
//
//Updates
//Author	:	Ross Conroy ross.conroy@tees.ac.uk
//Date		:	12/05/2014
//Included support for expanding no forgetting/DID IDs
//CPT's for next step exclude decision nodes,
//includes all previous time steps parents when adding parents to decision
//nodes,
//
//Author	:	Ross Conroy ross.conroy@tees.ac.uk
//Date		:	21/05/2014
//included support for expanding i-DIDs
//////////////////////////////////////////////////////////////////////////

ExpandTimeSteps::ExpandTimeSteps(void)
{
	GapBetweenSteps = 300;
}

Domain * ExpandTimeSteps::Expand(Domain * inputID, int noTimeSteps, IdType idTpye)
{
	int lastTimeStep = FindLastTimeStep(inputID);

	for(int i = 0; i < noTimeSteps; i++)
	{
		NodeList nlLastTimeStep = GetNodesForTimeStep(inputID, lastTimeStep);
		NodeList nlNextTimeStep = GetNodesForNextTimeStep(nlLastTimeStep, inputID);
		SetParentsForNextStep(nlLastTimeStep, idTpye);
		SetCPTsForNextStep(nlLastTimeStep, idTpye);

		lastTimeStep ++;
	}

	//cout << "save net" << endl;
	//inputID->saveAsNet("C:\\Nets\\output.net");

	return inputID;
}

//Uses the name of a node and extracts the node name and time step information
//and return the name for the next time step
string ExpandTimeSteps::GetNameForNextTimeStep(string previousName)
{
	int lastIndex = previousName.find_last_not_of("0123456789");
	string endNo = previousName.substr(lastIndex + 1);
	string nodeName = previousName.substr(0, lastIndex + 1);
	int step = atoi(endNo.c_str());
	step ++;

	stringstream ss;
	ss << nodeName << step;

	return ss.str();
}

void ExpandTimeSteps::SetCPTsForNextStep(NodeList previousNodes, IdType idTpye)
{
	for(NodeList::const_iterator it = previousNodes.begin(); it != previousNodes.end(); ++it)
	{
		Node * oldNode = *it;

		if(oldNode->getCategory() != H_CATEGORY_DECISION)
		{
			if(((idTpye == IDID) && !isAgentJActionNodeName(oldNode->getName())) || (idTpye != IDID))
			{
				Node * newNode = oldNode->getDomain()->getNodeByName(GetNameForNextTimeStep(oldNode->getName()));
				newNode->getTable()->setData(oldNode->getTable()->getData());
			}
		}				
	}
}

//loop through previous time step nodes and use GetNameForNextTimeStep
//to get the node names for the next time step to setup
void ExpandTimeSteps::SetParentsForNextStep(NodeList previousNodes, IdType idTpye)
{
	for(NodeList::const_iterator it = previousNodes.begin(); it != previousNodes.end(); ++it)
	{
		Node * oldNode = *it;

		if(oldNode->getCategory() == H_CATEGORY_UTILITY)
		{		
			UtilityNode * oldUtilNode = (UtilityNode *)oldNode;
			UtilityNode * newUtilNode = (UtilityNode *)oldNode->getDomain()->getNodeByName(GetNameForNextTimeStep(oldUtilNode->getName()));

			NodeList parents;
			parents = oldUtilNode->getParents();
			for(NodeList::reverse_iterator itp = parents.rbegin(); itp != parents.rend(); ++itp)
			{
				Node * parentNode = *itp;
				string newParentName = GetNameForNextTimeStep(parentNode->getName());
				newUtilNode->addParent((DiscreteNode *)oldNode->getDomain()->getNodeByName(newParentName));				
			}
		}
		else if(oldNode->getCategory() == H_CATEGORY_CHANCE)
		{		
			DiscreteChanceNode * oldChanceNode = (DiscreteChanceNode *)oldNode;
			DiscreteChanceNode * newChanceNode = (DiscreteChanceNode *)oldNode->getDomain()->getNodeByName(GetNameForNextTimeStep(oldChanceNode->getName()));

			NodeList parents;
			parents = oldChanceNode->getParents();
			for(NodeList::reverse_iterator itp = parents.rbegin(); itp != parents.rend(); ++itp)
			{
				Node * parentNode = *itp;
				string newParentName = GetNameForNextTimeStep(parentNode->getName());
				newChanceNode->addParent((DiscreteNode *)oldNode->getDomain()->getNodeByName(newParentName));

				//Update i-DID support, needs to expand agent j nodes in the same way as decision nodes with no forgetting
				if((idTpye == IDID) && (isAgentJActionNodeName(newChanceNode->getName())))
				{
					//If parent if from first time step then it needs adding too
					if(GetTimeStepFromNodeName(parentNode->getName()) == 0)
					{
						newChanceNode->addParent((DiscreteNode *) parentNode);
					}					
				}
			}			
		}
		else if(oldNode->getCategory() == H_CATEGORY_DECISION)
		{		
			DiscreteChanceNode * oldDecisionNodeNode = (DiscreteChanceNode *)oldNode;
			DiscreteChanceNode * newDecisionNodeNode = (DiscreteChanceNode *)oldNode->getDomain()->getNodeByName(GetNameForNextTimeStep(oldDecisionNodeNode->getName()));

			NodeList parents;
			parents = oldDecisionNodeNode->getParents();
			for(NodeList::reverse_iterator itp = parents.rbegin(); itp != parents.rend(); ++itp)
			{
				Node * parentNode = *itp;
				string newParentName = GetNameForNextTimeStep(parentNode->getName());
				newDecisionNodeNode->addParent((DiscreteNode *)oldNode->getDomain()->getNodeByName(newParentName));

				if((idTpye == DID) || (idTpye == NOFORGET) || (idTpye == IDID))
				{
					//If parent if from first time step then it needs adding too
					if(GetTimeStepFromNodeName(parentNode->getName()) == 0)
					{
						newDecisionNodeNode->addParent((DiscreteNode *) parentNode);
					}					
				}
			}			
		}
	}
}

//Checks to see if the node name contains the agent j prefix
bool ExpandTimeSteps::isAgentJActionNodeName(string nodeName)
{
	if (nodeName.find("aj") != std::string::npos) {
		return true;
	}
	else
	{
		return false;
	}
}

//extracts the time step from a node name
int ExpandTimeSteps::GetTimeStepFromNodeName(string inNodeName)
{
	int lastIndex = inNodeName.find_last_not_of("0123456789");
	string endNo = inNodeName.substr(lastIndex + 1);
	int step = atoi(endNo.c_str());

	return step;
}

//loops through each node and copies them,
//needs to determine prerequisites for each existing node and note them
//for the new nodes with time step incremented.
//parent node names and order are stored in the user data as an array of strings
NodeList ExpandTimeSteps::GetNodesForNextTimeStep(NodeList previousNodes, Domain * inputID)
{
	NodeList nextStepNodes;

	for(NodeList::const_iterator it = previousNodes.begin(); it != previousNodes.end(); ++it)
	{
		Node * oldNode = *it;
		Node * newNode = NULL;
		//NodeList parents;

		if(oldNode->getCategory() == H_CATEGORY_UTILITY)
		{
			UtilityNode * utilNode = new UtilityNode(inputID);

			UtilityNode * tempUtilNode = (UtilityNode *)oldNode;

			newNode = utilNode;
		}

		else if(oldNode->getCategory() == H_CATEGORY_CHANCE)
		{
			LabelledDCNode * chanceNode = new LabelledDCNode(inputID);

			DiscreteChanceNode * tempOldNode = (DiscreteChanceNode *)oldNode;
			chanceNode->setNumberOfStates(tempOldNode->getNumberOfStates());

			for(int i = 0; i < tempOldNode->getNumberOfStates(); i++)
			{
				chanceNode->setStateLabel(i, tempOldNode->getStateLabel(i));
			}

			newNode = chanceNode;
		}

		else if(oldNode->getCategory() == H_CATEGORY_DECISION)
		{
			LabelledDDNode * decisionNode = new LabelledDDNode(inputID);

			DiscreteDecisionNode * tempOldNode = (DiscreteDecisionNode *)oldNode;
			decisionNode->setNumberOfStates(tempOldNode->getNumberOfStates());

			for(int i = 0; i < tempOldNode->getNumberOfStates(); i++)
			{
				decisionNode->setStateLabel(i, tempOldNode->getStateLabel(i));
			}

			newNode = decisionNode;
		}

		newNode->setName(GetNameForNextTimeStep(oldNode->getName()));
		newNode->setPosition(oldNode->getPosition().first + GapBetweenSteps, oldNode->getPosition().second);

		nextStepNodes.push_back(newNode);
	}

	return nextStepNodes;
}

//Gets all the nodes for a time step
//TODO read string like next node code
NodeList ExpandTimeSteps::GetNodesForTimeStep(Domain * inputID, int timeStep)
{
	NodeList nl;

	NodeList nodes = inputID->getNodes();
	for(NodeList::const_iterator it = nodes.begin(); it != nodes.end(); ++it)
	{
		Node * node = *it;

		int lastIndex = node->getName().find_last_not_of("0123456789");
		string endNo = node->getName().substr(lastIndex + 1);
		int step = atoi(endNo.c_str());

		if(step == timeStep)
		{
			nl.push_back(node);
		}
	}

	return nl;
}

//This method loops through all nodes within the ID to find the one
//which ends in the highest number
int ExpandTimeSteps::FindLastTimeStep(Domain * inputID)
{
	int last = 0;

	NodeList nodes = inputID->getNodes();
	for(NodeList::const_iterator it = nodes.begin(); it != nodes.end(); ++it)
	{
		Node * node = *it;
		int lastIndex = node->getName().find_last_not_of("0123456789");
		string endNo = node->getName().substr(lastIndex + 1);
		int step = atoi(endNo.c_str());

		if(step > last)
		{
			last = step;
		}
	}

	return last;
}

int ExpandTimeSteps::FindFirstTimeStep(Domain * inputID)
{
	int first = 0;
	bool firstB = false;

	NodeList nodes = inputID->getNodes();
	for(NodeList::const_iterator it = nodes.begin(); it != nodes.end(); ++it)
	{
		Node * node = *it;
		int lastIndex = node->getName().find_last_not_of("0123456789");
		string endNo = node->getName().substr(lastIndex + 1);
		int step = atoi(endNo.c_str());

		if(step < first || !firstB)
		{
			first = step;
			firstB = !firstB;
		}
	}

	return first;
}
